#include "gemm.h"
#include "utils.h"
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <Accelerate/Accelerate.h>

void pm(int M, int N, float *A);

void gemm_bin(int M, int N, int K, float ALPHA,
              char *A, int lda,
              float *B, int ldb,
              float *C, int ldc) {
    int i, j, k;
    for (i = 0; i < M; ++i) {
        for (k = 0; k < K; ++k) {
            char A_PART = A[i * lda + k];
            if (A_PART) {
                for (j = 0; j < N; ++j) {
                    C[i * ldc + j] += B[k * ldb + j];
                }
            } else {
                for (j = 0; j < N; ++j) {
                    C[i * ldc + j] -= B[k * ldb + j];
                }
            }
        }
    }
}

float *random_matrix(int rows, int cols) {
    int i;
    float *m = calloc(rows * cols, sizeof(float));
    for (i = 0; i < rows * cols; ++i) {
        m[i] = (float) rand() / (float) RAND_MAX;
    }
    return m;
}

void time_random_matrix(int TA, int TB, int m, int k, int n) {
    float *a;
    if (!TA) a = random_matrix(m, k);
    else a = random_matrix(k, m);
    int lda = (!TA) ? k : m;
    float *b;
    if (!TB) b = random_matrix(k, n);
    else b = random_matrix(n, k);
    int ldb = (!TB) ? n : k;

    float *c = random_matrix(m, n);
    int i;
    clock_t start = clock(), end;
    for (i = 0; i < 10; ++i) {
        gemm_cpu(TA, TB, m, n, k, 1, a, lda, b, ldb, 1, c, n);
        // cblas_sgemm(CblasRowMajor, TA==0?CblasNoTrans:CblasTrans, TB==0?CblasNoTrans:CblasTrans,
        //         m, n, k, 1.0, a, lda, b, ldb, 1, c, n);
    }
    end = clock();
    printf("Matrix Multiplication %dx%d * %dx%d, TA=%d, TB=%d: %lf ms\n", m, k, k, n, TA, TB,
           (float) (end - start) / CLOCKS_PER_SEC);
    // pm(m,n,a);
    free(a);
    free(b);
    free(c);
}


void gemm(int TA, int TB, int M, int N, int K, float ALPHA,
          float *A, int lda,
          float *B, int ldb,
          float BETA,
          float *C, int ldc) {
    gemm_cpu(TA, TB, M, N, K, ALPHA, A, lda, B, ldb, BETA, C, ldc);
}

void gemm_nn(int M, int N, int K, float ALPHA,
             float *A, int lda,
             float *B, int ldb,
             float *C, int ldc) {
    int i, j, k;
#pragma omp parallel for
    for (i = 0; i < M; ++i) {
        for (k = 0; k < K; ++k) {
            register float A_PART = ALPHA * A[i * lda + k];
            for (j = 0; j < N; ++j) {
                C[i * ldc + j] += A_PART * B[k * ldb + j];
            }
        }
    }
}

void gemm_nt(int M, int N, int K, float ALPHA,
             float *A, int lda,
             float *B, int ldb,
             float *C, int ldc) {
    int i, j, k;
#pragma omp parallel for
    for (i = 0; i < M; ++i) {
        for (j = 0; j < N; ++j) {
            register float sum = 0;
            for (k = 0; k < K; ++k) {
                sum += ALPHA * A[i * lda + k] * B[j * ldb + k];
            }
            C[i * ldc + j] += sum;
        }
    }
}

void gemm_tn(int M, int N, int K, float ALPHA,
             float *A, int lda,
             float *B, int ldb,
             float *C, int ldc) {
    int i, j, k;
#pragma omp parallel for
    for (i = 0; i < M; ++i) {
        for (k = 0; k < K; ++k) {
            register float A_PART = ALPHA * A[k * lda + i];
            for (j = 0; j < N; ++j) {
                C[i * ldc + j] += A_PART * B[k * ldb + j];
            }
        }
    }
}

void gemm_tt(int M, int N, int K, float ALPHA,
             float *A, int lda,
             float *B, int ldb,
             float *C, int ldc) {
    int i, j, k;
#pragma omp parallel for
    for (i = 0; i < M; ++i) {
        for (j = 0; j < N; ++j) {
            register float sum = 0;
            for (k = 0; k < K; ++k) {
                sum += ALPHA * A[i + k * lda] * B[k + j * ldb];
            }
            C[i * ldc + j] += sum;
        }
    }
}


void gemm_cpu(int TA, int TB, int M, int N, int K, float ALPHA,
              float *A, int lda,
              float *B, int ldb,
              float BETA,
              float *C, int ldc) {
    /*
    printf("cpu: %d %d %d %d %d %f %d %d %f %d\n",TA, TB, M, N, K, ALPHA, lda, ldb, BETA, ldc);
     int i, j;
     for(i = 0; i < M; ++i){
         for(j = 0; j < N; ++j){
             C[i*ldc + j] *= BETA;
         }
     }
     if(!TA && !TB)
         gemm_nn(M, N, K, ALPHA,A,lda, B, ldb,C,ldc);
     else if(TA && !TB)
         gemm_tn(M, N, K, ALPHA,A,lda, B, ldb,C,ldc);
     else if(!TA && TB)
         gemm_nt(M, N, K, ALPHA,A,lda, B, ldb,C,ldc);
     else
         gemm_tt(M, N, K, ALPHA,A,lda, B, ldb,C,ldc);
         */
    cblas_sgemm(CblasRowMajor, TA == 0 ? CblasNoTrans : CblasTrans, TB == 0 ? CblasNoTrans : CblasTrans,
                M, N, K, ALPHA, A, lda, B, ldb, BETA, C, ldc);
}
